import json
import logging
from typing import Dict, List
from fastapi import WebSocket

logger = logging.getLogger(__name__)


class ConnectionManager:
    def __init__(self):
        # 设备连接管理
        self.active_device_connections: Dict[str, WebSocket] = {}
        # AI对话会话管理
        self.chat_sessions: Dict[str, List[Dict]] = {}

    async def connect_device(self, device_id: str, websocket: WebSocket):
        """连接设备"""
        await websocket.accept()
        self.active_device_connections[device_id] = websocket
        logger.info(f"设备 {device_id} 连接成功")

    def disconnect_device(self, device_id: str):
        """断开设备连接"""
        if device_id in self.active_device_connections:
            del self.active_device_connections[device_id]
        if device_id in self.chat_sessions:
            del self.chat_sessions[device_id]
        logger.info(f"设备 {device_id} 断开连接")

    async def send_personal_message(self, message: dict, device_id: str):
        """向指定设备发送消息"""
        if device_id in self.active_device_connections:
            try:
                await self.active_device_connections[device_id].send_text(json.dumps(message))
            except Exception as e:
                logger.error(f"向设备 {device_id} 发送消息失败: {str(e)}")

    async def broadcast(self, message: dict):
        """广播消息给所有设备"""
        for device_id, websocket in self.active_device_connections.items():
            try:
                await websocket.send_text(json.dumps(message))
            except Exception as e:
                logger.error(f"广播消息到设备 {device_id} 失败: {str(e)}")

    def get_chat_history(self, device_id: str) -> List[Dict]:
        """获取对话历史"""
        return self.chat_sessions.get(device_id, [])

    def update_chat_history(self, device_id: str, history: List[Dict]):
        """更新对话历史"""
        self.chat_sessions[device_id] = history

    def get_connected_devices(self) -> List[str]:
        """获取所有连接的设备ID"""
        return list(self.active_device_connections.keys())